"""
    create by IAmFiveHigh on 2024/8/21
"""


class AVLTree:
    class Node:
        def __init__(self, key, value, left=None, right=None):
            self.key = key
            self.value = value
            self.left = left
            self.right = right
            self.height = 1

    def __init__(self):
        self._root = None
        self._size = 0

    def get_size(self) -> int:
        return self._size

    def is_empty(self) -> bool:
        return self._size == 0

    def add(self, key, value):
        if self.is_empty():
            self._root = self.Node(key, value)
            self._size += 1
            return

        self._root = self._add_from_node(self._root, key, value)

    def _add_from_node(self, node, key, value):
        if node is None:
            self._size += 1
            return self.Node(key, value)

        if key < node.key:
            node.left = self._add_from_node(node.left, key, value)
        elif key > node.key:
            node.right = self._add_from_node(node.right, key, value)
        else:
            # node.key == key 修改value
            node.value = value

        # 更新当前node的高度
        node.height = max(self._get_node_height(node.left), self._get_node_height(node.right)) + 1

        # 更新当前node的平衡因子
        balance_factor = self._get_balance(node)

        # RR \方向链表
        if balance_factor < -1 and self._get_balance(node.right) <= 0:
            # 左旋
            node = self.left_rotate(node)
        # LL /方向链表
        elif balance_factor > 1 and self._get_balance(node.left) >= 0:
            # 右旋
            node = self.right_rotate(node)
        # RL >方向
        elif balance_factor < -1 and self._get_balance(node.right) > 0:
            node.right = self.right_rotate(node.right)
            node = self.left_rotate(node)
        # LR <方向
        elif balance_factor > 1 and self._get_balance(node.left) < 0:
            node.left = self.left_rotate(node.left)
            node = self.right_rotate(node)
        return node

    # 对node y 进行右旋转，返回旋转后的x
    #          y                         x
    #        /   \                     /   \
    #       x    T4    向右旋转       z       y
    #     /   \       -------->   /   \     /  \
    #    z    T3                 T1   T2   T3   T4
    #  /   \
    # T1   T2
    def right_rotate(self, y: Node):
        x = y.left
        t3 = x.right

        x.right = y
        y.left = t3

        # 更新x y结点的高度
        y.height = max(self._get_node_height(y.left), self._get_node_height(y.right)) + 1
        x.height = max(self._get_node_height(x.left), self._get_node_height(x.right)) + 1

        return x

    def left_rotate(self, y: Node):
        x = y.right
        t3 = x.left

        x.left = y
        y.right = t3

        y.height = max(self._get_node_height(y.left), self._get_node_height(y.right)) + 1
        x.height = max(self._get_node_height(x.left), self._get_node_height(x.right)) + 1

        return x

    # 返回最小值node也就是左下角
    def mini_node(self):
        if self.is_empty():
            raise Exception("Error: 无法在空树中获取最小node")
        return self._mini_node(self._root)

    def _mini_node(self, node):
        if node.left is not None:
            return self._mini_node(node.left)
        else:
            return node

    # 返回最大值node也就是右下角
    def max_node(self):
        if self.is_empty():
            raise Exception("Error: 无法在空树中获取最大node")
        return self._max_node(self._root)

    def _max_node(self, node):
        if node.right is not None:
            return self._max_node(node.right)
        else:
            return node

    def remove_key(self, key):
        if self.is_empty():
            raise Exception("Error: 无法在空树中删除元素")
        self._root = self._remove_key(self._root, key)

    def _remove_key(self, node, key):
        if node is None:
            return None

        ret_node = None
        if key < node.key:
            node.left = self._remove_key(node.left, key)
            ret_node = node
        elif key > node.key:
            node.right = self._remove_key(node.right, key)
            ret_node = node
        else:

            # key == node.key
            # 就是要删除当前node
            if node.left is None:
                right_node = node.right
                node.right = None
                self._size -= 1
                ret_node = right_node
            elif node.right is None:
                left_node = node.left
                node.left = None
                self._size -= 1
                ret_node = left_node
            else:
                # 既有左子树又有右子树
                # 找到node.right为根的最小node, 也就是当前node值大一点的下一个node
                successor = self._mini_node(node.right)
                successor.right = self._remove_key(node.right, successor.key)
                successor.left = node.left

                node.left = None
                node.right = None
                ret_node = successor

        if ret_node is None:
            return ret_node

        ret_node.height = max(self._get_node_height(ret_node.left), self._get_node_height(ret_node.right)) + 1

        balance_factor = self._get_balance(ret_node)
        # LL
        if balance_factor > 1 and self._get_balance(ret_node.left) >= 0:
            ret_node = self.right_rotate(ret_node)
        # LR
        elif balance_factor > 1 and self._get_balance(ret_node.left) < 0:
            ret_node.left = self.left_rotate(ret_node.left)
            ret_node = self.right_rotate(ret_node)
        # RR
        elif balance_factor < -1 and self._get_balance(ret_node.right) <= 0:
            ret_node = self.left_rotate(ret_node)
        # RL
        elif balance_factor < -1 and self._get_balance(ret_node.right) > 0:
            ret_node.right = self.right_rotate(ret_node.right)
            ret_node = self.left_rotate(ret_node)

        return ret_node

    # 判断是否为二分搜索树
    def is_bst(self):
        if self.is_empty():
            return False
        order_result = []
        self._in_order(self._root, lambda k, v: order_result.append(k))
        # 如果是二分搜索树，中序遍历的数组应该从小到大排序
        for index in range(1, len(order_result)):
            if order_result[index] < order_result[index - 1]:
                return False
        return True

    # 中序遍历 func为对遍历到的key value操作的方法
    def in_order(self, func):
        self._in_order(self._root, func)

    def _in_order(self, node, func):
        if node is None:
            return
        self._in_order(node.left, func)
        func(node.key, node.value)
        self._in_order(node.right, func)

    # 判断是否是平衡二叉树
    def is_balanced(self):
        return self._is_balanced(self._root)

    def _is_balanced(self, node):
        if node is None:
            return True
        balance_factor = self._get_balance(node)
        # 任意node的平衡因子不能大于1 否则就不算平衡二叉树
        if abs(balance_factor) > 1:
            return False

        return self._is_balanced(node.left) and self._is_balanced(node.right)

    # 获得node结点的高度
    @staticmethod
    def _get_node_height(node):
        if node is None:
            return 0
        return node.height

    def _get_balance(self, node):
        if node is None:
            return 0
        return self._get_node_height(node.left) - self._get_node_height(node.right)


if __name__ == '__main__':

    def test_tree(t, result_list, error_message):
        in_order_list = []
        t.in_order(lambda k, v: in_order_list.append(k))
        if in_order_list != result_list:
            print(error_message)


    temp = [5, 6, 1, 3, 4, 2]
    tree = AVLTree()
    for i in temp:
        tree.add(i, None)
    test_tree(tree, [1, 2, 3, 4, 5, 6], "二叉树添加元素方法有bug")
    tree.remove_key(3)
    test_tree(tree, [1, 2, 4, 5, 6], "二叉树删除元素方法有bug")

    if not tree.is_balanced():
        print("不是平衡二叉树")
